﻿// --------------------------------------------------------------------------------------------------------------------
// <copyright file="RedBlackTree.cs" company="Team Awesome">
//   Awesome
// </copyright>
// <summary>
//   Defines the RedBlackTree type.
// </summary>
// --------------------------------------------------------------------------------------------------------------------

namespace System.Collections.Generic
{
    using Diagnostics.Contracts;

    /// <summary>
    /// A red-black tree is a self-balancing binary search tree.
    /// The tree shall always hold the following properties:
    /// 1. A node is either red or black.
    /// 2. The root is black. (This rule is sometimes omitted from other definitions. Since the root can always be changed from red to black but not necessarily vice-versa this rule has little effect on analysis.)
    /// 3. All leaves are black.
    /// 4. Both children of every red node are black.
    /// 5. Every simple path from a given node to any of its descendant leaves contains the same number of black nodes.
    /// </summary>
    /// <typeparam name="T"> Type of items held in the tree nodes.
    /// </typeparam>
    public partial class RedBlackTree<T> where T : class
    {
        /// <summary>
        /// The comparer for the type T to be used.
        /// </summary>
        private readonly IComparer<T> _comparer;

        /// <summary>
        /// The tree's root node.
        /// </summary>
        private Node _root;

        /// <summary>
        /// Initializes a new instance of the <see cref="RedBlackTree{T}"/> class. 
        /// Create a red-black tree collection with natural comparer.
        /// </summary>
        public RedBlackTree() : this(Comparer<T>.Default)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="RedBlackTree{T}"/> class. 
        /// Create a red-black tree collection with an external comparer. 
        /// </summary>
        /// <param name="comparer">
        /// The external comparer
        /// </param>
        public RedBlackTree(IComparer<T> comparer)
        {
            Contract.Requires<NullReferenceException>(comparer != null, "Item comparer cannot be null");

            _comparer = comparer;
        }

        /// <summary>
        /// Method for inserting a new item into the tree. This will automatically refactor the tree in order
        /// for the invariances to be hold.
        /// </summary>
        /// <param name="item">
        /// The item.
        /// </param>
        public void Add(T item)
        {
            #region Code Contracs: Pre - and Post conditions

            Contract.Requires(item != null);

            #endregion

            _root = Add(_root, item);
            _root.IsRed = false;
        }

        /// <summary>
        /// Takes a node and tests if it should move further down the tree or if it should insert right away.
        /// </summary>
        /// <param name="node">
        /// The node.
        /// </param>
        /// <param name="item">
        /// The item.
        /// </param>
        /// <returns>
        /// The inserted node.
        /// </returns>
        private Node Add(Node node, T item)
        {
            #region Code Contracs: Pre - and Post conditions

            Contract.Requires(item != null);

            #endregion

            if (node == null) return new Node(item);

            var cmp = _comparer.Compare(item, node.Item);
            if (cmp == 0)
            {
                return node;
            }

            if (cmp < 0)
                node.Left = Add(node.Left, item);
            else
                node.Right = Add(node.Right, item);

            return FixSetup(node);
        }

        /// <summary>
        /// Validates the tree around the node and fixes everything not in place.
        /// </summary>
        /// <param name="node">
        /// The node.
        /// </param>
        /// <returns>
        /// The node corresponding to the same place in the tree.
        /// </returns>
        private static Node FixSetup(Node node)
        {
            #region Code Contracs: Pre - and Post conditions

            Contract.Requires(node != null);
            Contract.Ensures(Contract.Result<Node>() != null);
            Contract.Ensures(!Contract.Result<Node>().Right.IsRed);

            #endregion

            if (IsColourRed(node.Right))
                node = RotateLeft(node);
            if (IsColourRed(node.Left) && IsColourRed(node.Left.Left))
                node = RotateRight(node);
            if (IsColourRed(node.Left) && IsColourRed(node.Right))
                FlipColors(node);
            return node;
        }

        /// <summary>
        /// Flips the colors on the node and children of current node.
        /// </summary>
        /// <param name="node">
        /// The node.
        /// </param>
        private static void FlipColors(Node node)
        {
            #region Code Contracs: Pre - and Post conditions

            Contract.Requires(node.Left != null);
            Contract.Requires(node.Right != null);

            Contract.Ensures(Contract.OldValue(node).IsRed == !node.IsRed);
            Contract.Ensures(Contract.OldValue(node).Left.IsRed == !node.Left.IsRed);
            Contract.Ensures(Contract.OldValue(node).Right.IsRed == !node.Right.IsRed);

            #endregion

            node.IsRed = !node.IsRed;
            node.Left.IsRed = !node.Left.IsRed;
            node.Right.IsRed = !node.Right.IsRed;
        }

        /// <summary>
        /// If the weight of the tree (or subtree) is on the left side, then we want to turn the tree to the right.
        /// The roots left side node then becomes root itself and the old root takes the new roots right node.
        /// </summary>
        /// <param name="subRoot">
        /// The sub Root.
        /// </param>
        /// <returns>
        /// The new subroot.
        /// </returns>
        private static Node RotateRight(Node subRoot)
        {
            #region Code Contracs: Pre - and Post conditions

            // We dont have any reason to turn the tree right if there are nothing on the left side.
            Contract.Requires(subRoot != null);
            Contract.Requires(subRoot.Left != null);

            // subroot doesn't change node
            Contract.Ensures(Contract.Result<Node>().Right.Equals(subRoot));
            Contract.Ensures(Contract.OldValue(subRoot).Left.Equals(Contract.Result<Node>()));
            Contract.Ensures(Contract.OldValue(subRoot).Right.Equals(subRoot.Right));
            Contract.Ensures(Contract.OldValue(subRoot).Left.Right.Equals(subRoot.Left));
            Contract.Ensures(Contract.OldValue(subRoot).Left.Left.Equals(Contract.Result<Node>().Left));

            #endregion

            var node = subRoot.Left;
            subRoot.Left = node.Right;
            node.Right = subRoot;
            node.IsRed = subRoot.IsRed;
            subRoot.IsRed = true;
            return node;
        }

        /// <summary>
        /// If the weight of the tree (or subtree) is on the right side, then we want to turn the tree to the left.
        /// The roots right side node then becomes root itself and the old root takes the new roots left node.
        /// </summary>
        /// <param name="subRoot">
        /// The sub Root.
        /// </param>
        /// <returns>
        /// The new subroot.
        /// </returns>
        private static Node RotateLeft(Node subRoot)
        {
            #region Code Contracs: Pre - and Post conditions

            // We dont have any reason to turn the tree left if there are nothing on the right side.
            Contract.Requires(subRoot != null);
            Contract.Requires(subRoot.Right != null);
            
            // subroot doesn't change node
            Contract.Ensures(Contract.Result<Node>().Left.Equals(subRoot));
            Contract.Ensures(Contract.OldValue(subRoot).Right.Equals(Contract.Result<Node>()));

            Contract.Ensures(Contract.OldValue(subRoot).Left == subRoot.Left);
            Contract.Ensures(Contract.OldValue(subRoot).Right.Left.Equals(subRoot.Right));
            Contract.Ensures(Contract.OldValue(subRoot).Right.Right.Equals(Contract.Result<Node>().Right));

            #endregion

            var node = subRoot.Right;
            subRoot.Right = node.Left;
            node.Right = subRoot;
            node.IsRed = subRoot.IsRed;
            subRoot.IsRed = true;
            return node;
        }

        #region ContractInvariantMethod

        // ReSharper disable UnusedMember.Local

        /// <summary>
        /// This method is used by code contracts to set all the invariants for the class.
        /// </summary>
        [ContractInvariantMethod]
        private void ObjectInvariant()
        {    
            // In our implementation every leaf is the value null and represents the black nil.

            //// 1. A node is either red or black.
            //// 2. The root is black. (This rule is sometimes omitted from other definitions. Since the root can always be changed from red to black but not necessarily vice-versa this rule has little effect on analysis.)
            //// 3. All leaves are black.
            //// 5. Every simple path from a given node to any of its descendant leaves contains the same number of black nodes.

            // The root is black.
            Contract.Invariant(this._root == null || !this._root.IsRed);

            // No path from the root to the bottom contains two consecutive red links. 
            ////Contract.Invariant(Contract.ForAll(this.GetEnumerator(), (Node x) => x.IsBlack || (!x.IsBlack && ((x.Left == null || x.Left.IsBlack) && (x.Right == null || x.Right.IsBlack)))));

            Contract.Invariant(EqualAmountOfBlackNodes());

            // There can't be any right leaning 3 nodes or 4 nodes in a 2-3 LLRB tree.
            Contract.Invariant(No4NodesOrRl3Nodes());
        }

        /// <summary>
        /// Gets the number of black children or returns -1 if they haven't been equal from bottom up.
        /// </summary>
        /// <param name="node">
        /// The node.
        /// </param>
        /// <returns>
        /// Number of black children and -1 if they havent been equal from bottom up.
        /// </returns>
        [Pure]
        private static int AmountOfBlackChildren(Node node)
        {
            if (node == null)
                return 0;

            var leftBlackChildren = AmountOfBlackChildren(node.Left);
            var rightBlackChildren = AmountOfBlackChildren(node.Right);

            if (leftBlackChildren == -1 || rightBlackChildren == -1 || leftBlackChildren != rightBlackChildren)
                return -1;

            var blackChildren = leftBlackChildren + rightBlackChildren;

            return blackChildren + (!node.IsRed ? 1 : 0);
        }

        /// <summary>
        /// Checks if there is the same amount of black notes in the left and right subtree from every subtree in the tree.
        /// </summary>
        /// <returns>
        /// A boolean stating if there is the same amount of black notes in the left and right subtree from every subtree in the tree.
        /// </returns>
        [Pure]
        private bool EqualAmountOfBlackNodes()
        {
            return -1 != AmountOfBlackChildren(_root);
        }

        /// <summary>
        /// Checks if the tree is without 4 nodes or right leaning 3 nodes. That is, if any node has a red child to the right.
        /// </summary>
        /// <returns>
        /// Returns false if a 4 node or right leaning 3 node is found. Else true.
        /// </returns>
        [Pure]
        private bool No4NodesOrRl3Nodes()
        {
            return No4NodesOrRl3Nodes(_root);
        }

        /// <summary>
        /// Checks if the tree is without 4 nodes or right leaning 3 nodes. That is, if any node has a red child to the right.
        /// </summary>
        /// <param name="h">
        /// The node which to check if its subtrees contains 4 nodes or right leaning 3 nodes.
        /// </param>
        /// <returns>
        /// Returns false if a 4 node or right leaning 3 node is found. Else true.
        /// </returns>
        [Pure]
        private static bool No4NodesOrRl3Nodes(Node h)
        {
            if (h == null)
            {
                return true;
            }

            // Is right leaning 3 node or 4 node
            if (h.Right != null && h.Right.IsRed)
            {
                return false;
            }

            // There are two red nodes in a row
            if (h.IsRed && h.Left != null && h.Left.IsRed)
            {
                return false;
            }

            return No4NodesOrRl3Nodes(h.Left) && No4NodesOrRl3Nodes(h.Right);
        }

        // ReSharper restore UnusedMember.Local
        #endregion
    }
}
